#ifndef BOOTSTRAPPING_TRANSFORM_H
#define BOOTSTRAPPING_TRANSFORM_H
#include <vector>
#include <cassert>
#include <unordered_map>
#include <iostream>
#include "polyarith.h"
#include "slots.h"

template <class T> constexpr inline std::size_t HashCombine(T const & v, std::size_t const seed = {}) noexcept
{
    return seed ^ (std::hash<T>()(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2));
}

template <typename... Args> struct std::hash<std::tuple<Args...>> {
    std::size_t operator () (const std::tuple<Args...> &tuple) const noexcept
    {
        return std::apply(
            [](const auto &... args) {
                auto seed = static_cast<std::size_t>(672807365);
                ((seed = HashCombine(args, seed)), ...);
                return seed;
            },
            tuple);
    }
};
/* *
 * @class NegacyclicPowerTable
 * @brief 用于预先计算并存储给定元素的所有幂，以便快速访问。
 *
 * `NegacyclicPowerTable` 类提供了一个高效的方式来存储并快速访问给定生成元的所有幂。
 * 通过预先计算生成元的幂并将其存储在一个数组中，可以避免在后续计算中反复计算幂的值，提升性能。
 */
// 预先计算给定元素的所有幂并存储它们以便快速访问。
class NegacyclicPowerTable {
    SlotRing::SubringView ring;
    uint64_t N;
    poly generator;
    std::vector<poly> content;

public:
    NegacyclicPowerTable(const SlotRing::SubringView &ring, poly generator, uint64_t halfOrderGenerator);
    NegacyclicPowerTable(const NegacyclicPowerTable &) = default;
    NegacyclicPowerTable(NegacyclicPowerTable &&) = default;
    ~NegacyclicPowerTable() = default;

    poly operator[](int64_t i) const;
};

class CompiledSubringLinearTransform;


/* *
 * @class CompiledLinearTransform
 * @brief 存储满足Galois自同构的线性变换。
 *
 * `CompiledLinearTransform` 类提供了一种方式来表示和存储线性变换，特别是与 Galois 自同构相关的变换。
 * 主要方法包括计算自同构、变换系数、编译 Frobenius 变换、以及支持通过“小步-大步”算法处理变换等。
 */
class CompiledLinearTransform {
    std::shared_ptr<const SlotRing> slotRing;

    // 存储与大步小步相关的自同构系数
    std::vector<poly> coefficients;

    CompiledLinearTransform(const std::shared_ptr<const SlotRing> &slotRing, std::vector<poly> coefficients);

    // 计算自同构,index为下标
    SlotRing::RawAuto Automorphism(size_t index) const;

    // 计算自同构,从 from下标自同构到to下标
    SlotRing::RawAuto DifferenceAutomorphism(size_t from, size_t to) const;

    SlotRing::RawAuto ReverseAutomorphism(size_t index) const;

    uint64_t G1SubgroupOrder() const;

    uint64_t G2SubgroupOrder() const;

    // 改变线性变换的方式f->f.
    void AddScaledTransform(const poly &scaling, const SlotRing::Rotation &rotation,
        const SlotRing::Frobenius &frobenius);

    size_t BabystepAutomorphismCount() const;
    size_t GiantstepAutomorphismCount() const;


    //     应用变换的小步-大步方法的副作用是在求值过程中对系数应用了大步自同构。该函数将返回系数变换
    void FixCoefficientShift();

    static std::vector<poly> CompileFrobenius(
        const std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &sparseTransformMatrix,
        const SlotRing::SubringView &ring, uint64_t p, size_t N, const NegacyclicPowerTable &powerTable);

public:
    CompiledLinearTransform(const CompiledLinearTransform &) = default;
    CompiledLinearTransform(CompiledLinearTransform &&) = default;
    ~CompiledLinearTransform() = default;

    template <typename T>
    static CompiledLinearTransform CompileSlotBasis(std::shared_ptr<const SlotRing> slotRing,
        T sparse_transform_matrix_per_slot, bool use_g2 = true);

    static CompiledLinearTransform ScalarSlotsToFirstCoefficients(std::shared_ptr<const SlotRing> slotRing);
    static CompiledLinearTransform FirstCoefficientsToScalarSlots(std::shared_ptr<const SlotRing> slotRing);


    CompiledSubringLinearTransform InRing() &&;

    friend class CompiledSubringLinearTransform;
};

/* *
 * @class CompiledSubringLinearTransform
 * @brief 表示可预测地只作用于子向量的线性变换。
 */
class CompiledSubringLinearTransform {
    CompiledLinearTransform subringTransform;
    std::shared_ptr<const SlotRing> slotRing;
    mutable std::vector<Plaintext> coefficientsPlain = {};

    SlotRing::RawAuto Automorphism(size_t index) const;
    SlotRing::RawAuto DifferenceAutomorphism(size_t from, size_t to) const;
    size_t BabystepAutomorphismCount() const;
    size_t GiantstepAutomorphismCount() const;

public:
    CompiledSubringLinearTransform(CompiledLinearTransform &&transform, std::shared_ptr<const SlotRing> newRing);
    CompiledSubringLinearTransform(const CompiledSubringLinearTransform &) = default;
    CompiledSubringLinearTransform(CompiledSubringLinearTransform &&) = default;
    ~CompiledSubringLinearTransform() = default;

    static CompiledSubringLinearTransform SlotsToCoeffs(std::shared_ptr<const SlotRing> slotRing);
    static CompiledSubringLinearTransform CoeffsToSlots(std::shared_ptr<const SlotRing> slotRing);

    poly operator () (const poly &x) const;
    // 线性变换主函数
    void ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk, Ciphertext &result) const;

    // 返回一组galois元素，使得后续提前生成后续进行的galois变换对应的密钥
    std::vector<uint32_t> GaloisElements() const;

    CompiledLinearTransform &&transform() &&;
};

void TestFirstCoeffsToScalarSlots();
void TestCompileSlotBasis();
void TestApplyCiphertext();
void TestApplyCiphertextSubring();

/* *
 * @brief 编译给定稀疏变换矩阵的线性变换。
 *
 * 使用给定的稀疏变换矩阵 `sparse_transform_matrix_per_slot` 为给定的 `SlotRing` 编译线性变换。
 * 在计算过程中，首先根据 `slotRing` 获取槽位的旋转操作，然后基于稀疏变换矩阵计算变换系数，并生成 Frobenius
 * 变换，最后将结果按需要的方式进行组合。
 *
 * @tparam T 稀疏变换矩阵的类型
 * @param slotRing 指向 `SlotRing` 的共享指针，表示要进行变换的环。
 * @param sparse_transform_matrix_per_slot 用于变换计算的稀疏变换矩阵函数。
 * @param use_g2 是否使用 G2 子群的标志。
 * @return 编译后的线性变换 `CompiledLinearTransform` 对象。
 */
template <typename T>
inline CompiledLinearTransform CompiledLinearTransform::CompileSlotBasis(std::shared_ptr<const SlotRing> slotRing,
    T sparse_transform_matrix_per_slot, bool use_g2)
{
    size_t blockSize = slotRing->SlotGroupLen();
    if (!use_g2) {
        blockSize = blockSize / 2;
    }
    std::vector<SlotRing::Rotation> rotations;
    for (size_t i = 0; i < blockSize; ++i) {
        rotations.push_back(slotRing->BlockRotate(i, blockSize));
    }
    std::unique_ptr<SlotRing::Rotation> laneSwitch;
    if (!use_g2) {
        laneSwitch = std::make_unique<SlotRing::Rotation>(slotRing->Rotate(blockSize));
    }

    const size_t d = slotRing->SlotRank();
    std::unordered_map<std::tuple<size_t, size_t>, uint64_t> slotwiseMatrix;
    NegacyclicPowerTable powerTable(slotRing->R(), slotRing->FromSlotValue({ 0, 1 }, 0), slotRing->N());

    std::vector<poly> coefficients;
    coefficients.resize(blockSize * d);
    for (auto &c : coefficients) {
        c.resize(slotRing->N());
    }
    CompiledLinearTransform result(slotRing, std::move(coefficients));

    for (size_t s = 0; s < blockSize; ++s) {
        for (size_t j = 0; j < blockSize; ++j) {
            const size_t blockRow = j;
            const size_t blockCol = (j + blockSize - s) % blockSize;
            slotwiseMatrix.clear();
            sparse_transform_matrix_per_slot(slotwiseMatrix, blockRow, blockCol);
            if (slotwiseMatrix.size() == 0) {
                continue;
            }
            std::vector<poly> frobeniusForm = CompiledLinearTransform::CompileFrobenius(slotwiseMatrix,
                slotRing->slot(), slotRing->Prime(), slotRing->N(), powerTable);
            for (size_t l = 0; l < d; ++l) {
                poly coeff = rotations[j](frobeniusForm[l]);
                result.AddScaledTransform(coeff, rotations[s], slotRing->Frobenius_(l));
                if (!use_g2) {
                    coeff = (*laneSwitch)(coeff);
                    result.AddScaledTransform(coeff, rotations[s], slotRing->Frobenius_(l));
                }
            }
        }
    }
    result.FixCoefficientShift();
    return result;
}
#endif